#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
统计CtripHotelBaseInfo集合中信息缺失的酒店数量
"""

import sys
import io

# 修复Windows控制台中文乱码问题
if sys.platform == 'win32':
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')
    sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8')

from pymongo import MongoClient
from config import MONGODB_CONFIG
from datetime import datetime


def connect_db():
    """连接数据库"""
    client = MongoClient(
        host=MONGODB_CONFIG['host'],
        port=MONGODB_CONFIG['port'],
        username=MONGODB_CONFIG['username'],
        password=MONGODB_CONFIG['password']
    )
    db = client[MONGODB_CONFIG['db_name']]
    hotel_collection = db[MONGODB_CONFIG['collection_name']]
    return client, hotel_collection


def count_total(collection):
    """统计总数"""
    total = collection.estimated_document_count()
    print(f"\n酒店总数: {total:,}")
    return total


def has_english(text):
    """检查文本中是否包含英文字符"""
    if not text:
        return False
    return any(c.isalpha() and ord(c) < 128 for c in str(text))


def has_chinese(text):
    """检查文本中是否包含中文字符"""
    if not text:
        return False
    return any('\u4e00' <= c <= '\u9fff' for c in str(text))


def check_missing_rules(doc):
    """
    检查酒店信息是否缺失
    返回: (是否缺失, 缺失原因列表)
    
    根据规则检查：
    1. 酒店名称没有英文
    2. 酒店无主图
    3. 酒店缺失地址
    4. 酒店缺失经纬度信息
    5. 酒店类型缺失
    6. 缺失品牌名称
    """
    missing_reasons = []
    
    # 获取基础信息
    hotel_base_info = doc.get('hotelBaseInfo', {})
    
    # 规则1: 酒店名称没有英文
    hotel_name = hotel_base_info.get('hotelName', '')
    hotel_en_name = hotel_base_info.get('hotelEnName', '')
    
    if not has_english(hotel_name) and not has_english(hotel_en_name):
        missing_reasons.append('酒店名称没有英文')
    
    # 规则3: 酒店无主图
    hotel_picture_info = hotel_base_info.get('hotelPictureInfo', {})
    # 检查主图是否存在且非空
    if not hotel_picture_info or len(hotel_picture_info) == 0:
        missing_reasons.append('酒店无主图')
    else:
        # 检查是否有有效的图片数据
        has_valid_picture = False
        for key, value in hotel_picture_info.items():
            if value and value != '' and value != {}:
                has_valid_picture = True
                break
        if not has_valid_picture:
            missing_reasons.append('酒店无主图')
    
    # 规则4: 酒店缺失地址
    position_info = hotel_base_info.get('hotelPositionInfo', {})
    hotel_address = position_info.get('hotelAddress', '').strip()
    
    if not hotel_address or hotel_address == '':
        missing_reasons.append('酒店缺失地址')
    
    # 规则5: 酒店缺失经纬度信息
    coordinate_info = position_info.get('coordinateInfo', {})
    g_lon = coordinate_info.get('gLon', -1)
    g_lat = coordinate_info.get('gLat', -1)
    
    # 检查经纬度是否有效（不为-1且不为0）
    if g_lon in [-1, 0] or g_lat in [-1, 0]:
        missing_reasons.append('酒店缺失经纬度信息')
    
    # 规则6: 酒店类型缺失或无法匹配
    category_info = hotel_base_info.get('hotelCategoryInfo', {})
    category_id = category_info.get('categoryId', 0)
    category_name = category_info.get('categoryName', '')
    
    if not category_id or category_id == 0 or not category_name:
        missing_reasons.append('酒店类型缺失')
    
    # 规则7: 缺失品牌名称
    brand_info = hotel_base_info.get('hotelBrandInfo', {})
    brand_id = brand_info.get('brandId', 0)
    
    if not brand_id or brand_id == 0:
        missing_reasons.append('缺失品牌')
    
    is_missing = len(missing_reasons) > 0
    return is_missing, missing_reasons


def count_missing_hotels(hotel_collection, sample_size=None):
    """
    统计缺失信息的酒店
    
    Args:
        hotel_collection: 酒店集合
        sample_size: 抽样数量，None表示全量统计
    """
    print("\n" + "=" * 70)
    print("开始统计信息缺失的酒店...")
    print("=" * 70)
    
    total_count = 0
    missing_count = 0
    missing_details = {}  # 记录各种缺失原因的数量
    
    # 构建查询
    cursor = hotel_collection.find({})
    if sample_size:
        cursor = cursor.limit(sample_size)
        print(f"\n抽样模式: 检查前 {sample_size:,} 条数据")
    else:
        print(f"\n全量模式: 检查所有数据")
    
    cursor = cursor.batch_size(1000)
    
    print("处理酒店数据中...")
    import sys
    start_time = datetime.now()
    
    for doc in cursor:
        total_count += 1
        
        # 检查是否缺失
        is_missing, reasons = check_missing_rules(doc)
        
        if is_missing:
            missing_count += 1
            
            # 统计各种缺失原因
            for reason in reasons:
                missing_details[reason] = missing_details.get(reason, 0) + 1
        
        # 显示进度（每5000条显示一次）
        if total_count % 5000 == 0:
            elapsed = (datetime.now() - start_time).total_seconds()
            speed = total_count / elapsed if elapsed > 0 else 0
            progress = total_count / 2554437 * 100 if not sample_size else total_count / sample_size * 100
            print(f"  已处理: {total_count:,} 条 ({progress:.1f}%) | 缺失: {missing_count:,} | 速度: {speed:.0f} 条/秒")
            sys.stdout.flush()  # 强制刷新输出
    
    end_time = datetime.now()
    elapsed = (end_time - start_time).total_seconds()
    
    # 显示结果
    print("\n" + "=" * 70)
    print("统计完成！")
    print("=" * 70)
    print(f"处理总数: {total_count:,}")
    print(f"信息缺失酒店数: {missing_count:,}")
    print(f"信息完整酒店数: {total_count - missing_count:,}")
    print(f"缺失率: {missing_count / total_count * 100:.2f}%")
    print(f"耗时: {elapsed:.2f} 秒")
    
    if missing_details:
        print(f"\n缺失原因统计:")
        for reason, count in sorted(missing_details.items(), key=lambda x: x[1], reverse=True):
            print(f"  - {reason}: {count:,} 条 ({count / total_count * 100:.2f}%)")
    
    result = {
        'total_count': total_count,
        'missing_count': missing_count,
        'complete_count': total_count - missing_count,
        'missing_rate': missing_count / total_count * 100 if total_count > 0 else 0,
        'missing_details': missing_details,
        'elapsed_seconds': elapsed
    }
    
    # 保存结果到文件
    save_result_to_file(result, sample_size)
    
    return result


def save_result_to_file(result, sample_size=None):
    """保存统计结果到文件"""
    import os
    
    # 确保output目录存在
    os.makedirs('output', exist_ok=True)
    
    # 生成文件名
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    if sample_size:
        filename = f"output/酒店信息缺失结果_抽样{sample_size}_{timestamp}.txt"
    else:
        filename = f"output/酒店信息缺失结果_{timestamp}.txt"
    
    with open(filename, 'w', encoding='utf-8') as f:
        f.write("=" * 70 + "\n")
        f.write("酒店信息缺失统计报告\n")
        f.write("=" * 70 + "\n")
        f.write(f"生成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write(f"统计模式: {'抽样' if sample_size else '全量'}\n")
        if sample_size:
            f.write(f"抽样数量: {sample_size:,}\n")
        f.write("\n")
        
        f.write("-" * 70 + "\n")
        f.write("总体统计\n")
        f.write("-" * 70 + "\n")
        f.write(f"处理总数: {result['total_count']:,}\n")
        f.write(f"信息缺失酒店数: {result['missing_count']:,}\n")
        f.write(f"信息完整酒店数: {result['complete_count']:,}\n")
        f.write(f"缺失率: {result['missing_rate']:.2f}%\n")
        f.write(f"处理耗时: {result['elapsed_seconds']:.2f} 秒\n")
        f.write("\n")
        
        if result['missing_details']:
            f.write("-" * 70 + "\n")
            f.write("缺失原因详细统计\n")
            f.write("-" * 70 + "\n")
            f.write(f"{'序号':<6}{'缺失原因':<30}{'数量':<15}{'占比':<10}\n")
            f.write("-" * 70 + "\n")
            
            total = result['total_count']
            for idx, (reason, count) in enumerate(
                sorted(result['missing_details'].items(), key=lambda x: x[1], reverse=True), 1
            ):
                percentage = count / total * 100
                f.write(f"{idx:<6}{reason:<30}{count:<15,}{percentage:.2f}%\n")
        
        f.write("\n")
        f.write("=" * 70 + "\n")
        f.write("报告结束\n")
        f.write("=" * 70 + "\n")
    
    print(f"\n[OK] 统计结果已保存到: {filename}")
    return filename


def main():
    """主函数"""
    import sys
    
    # 连接数据库
    print("\n正在连接MongoDB...")
    print(f"数据库: {MONGODB_CONFIG['db_name']}")
    print(f"集合: {MONGODB_CONFIG['collection_name']}")
    
    client, hotel_collection = connect_db()
    print("[OK] 连接成功")
    
    # 统计总数
    count_total(hotel_collection)
    
    # 检查命令行参数
    sample_size = None
    if len(sys.argv) > 1:
        try:
            sample_size = int(sys.argv[1])
        except ValueError:
            print(f"错误: 参数必须是数字")
            return
    
    # 统计缺失酒店
    result = count_missing_hotels(hotel_collection, sample_size)
    
    # 关闭连接
    client.close()
    
    print("\n" + "=" * 70)
    print("\n使用说明:")
    print("  python count_missing_hotels.py           # 全量统计")
    print("  python count_missing_hotels.py 1000      # 抽样统计（前1000条）")
    print("=" * 70 + "\n")


if __name__ == '__main__':
    main()

